/**
 * @file ymodem.c
 * @author QuanShen (quan.shen@prime-semi.com)
 * @brief 
 * @version 0.1
 * @date 2025-07-29
 * 
 * @copyright Copyright (c) 2025
 * 
 */
#include "ymodem.h"
#include <stdint.h>
#include <stddef.h>
#include <string.h>
#include "flash_if.h"
#include "common.h"

// extern uint8_t file_name[];
uint8_t file_name[FILE_NAME_LENGTH];    // 文件名

static uint16_t UpdateCRC16(uint16_t crc_in, uint8_t byte);
static uint16_t Cal_CRC16(const uint8_t *data, uint32_t length);
static uint8_t CalChecksum(const uint8_t *data, uint32_t length);
static void Ymodem_SendPacket(uint8_t *data, uint16_t length);
static void Ymodem_PreparePacket(uint8_t *source_buff, uint8_t *data,
						  uint8_t pk_no, uint16_t size_blk);
static void Ymodem_PrepareIntialPacket(uint8_t *data, const uint8_t* f_name, uint16_t *length);

/**
 * @brief 接收字符
 * 
 * @param ch 字符
 * @param timeout 超时时间 
 * @return int32_t 0 : 接收成功, -1 : 超时或错误
 */
static int32_t Receive_Byte(uint8_t *ch, uint32_t timeout)
{
	while (timeout-- > 0) 
	{
		if (1 == SerialKeyPressed(ch))
		{
			return 0;
		}
	}
	return -1; // Timeout or error
}

/**
 * @brief 发送一个字节
 * 
 * @param ch 要发送的字节
 * @return uint32_t 0 : 发送成功
 */
static uint32_t Send_Byte(uint8_t ch)
{
	SerialPutChar(ch);
	return 0; // Assuming send is always successful
}

/**
 * @brief 更新 CRC-16 校验和
 * 
 * @param crc_in 当前 CRC-16 值
 * @param byte 要更新的字节
 * @return uint16_t 更新后的 CRC-16 值
 */
static uint16_t UpdateCRC16(uint16_t crc_in, uint8_t byte)
{
	uint32_t crc = crc_in;
	uint32_t in = byte | 0x100;

	do
	{
		crc <<= 1;
		in <<= 1;

		if (in & 0x100)
		{
			crc++;
		}

		if (crc & 0x10000)
		{
			crc ^= 0x1021;
		}
	} while (!(in & 0x10000));
	
	return (uint16_t)(crc & 0xFFFF);
}

/**
 * @brief 计算 CRC-16 校验和
 * 
 * @param data 数据缓冲区
 * @param length 数据长度
 * @return uint16_t 计算得到的 CRC-16 值
 */
static uint16_t Cal_CRC16(const uint8_t *data, uint32_t length)
{
	uint16_t crc = 0; // Initial CRC value
	const uint8_t *end = data + length;
	
	while (data < end)
	{
		crc = UpdateCRC16(crc, *data++);
	}
	crc = UpdateCRC16(crc, 0); // Finalize CRC with zero byte
	crc = UpdateCRC16(crc, 0); // Finalize CRC with another zero
	
	return (uint16_t)(crc & 0xFFFF);
}

/**
 * @brief 计算数据的校验和
 * 
 * @param data 数据缓冲区
 * @param length 数据长度
 * @return uint8_t 计算得到的校验和
 */
static uint8_t CalChecksum(const uint8_t *data, uint32_t length)
{
	uint32_t checksum = 0;
	
	const uint8_t *end = data + length;

	while (data < end)
	{
		checksum += *data++;
	}
	return (uint8_t)(checksum & 0xFF);
}

/**
 * @brief 接收一个数据包
 * 
 * @param data 接收的数据包缓冲区
 * @param length 数据包长度
 * @param timeout 超时时间
 * @return int32_t 0 : 成功, -1 : 错误或超时, 1 : 用户中止
 */
static int32_t Receive_Packet(uint8_t *data, int32_t *length, uint32_t timeout)
{
	uint16_t i, packet_size, computed_crc;
	uint8_t c;
	*length = 0;

	if (Receive_Byte(&c, timeout) < 0)
	{
		return TIMEOUT_OR_PACKET_ERROR; // Timeout or error
	}

	switch (c)
	{
		case SOH: // 128 Byte packet
			packet_size = PACKET_SIZE;
			break;
		case STX: // 1K Byte packet
			packet_size = PACKET_1K_SIZE;
			break;
		case EOT: // End of Transmission
			return PACKET_RECEIVED;
		case CAN: // Cancel
			if ((Receive_Byte(&c, timeout) == 0 && c == CAN))
			{
				*length = -1; // Indicate cancel
				return PACKET_RECEIVED;
			}
			else
			{
				return TIMEOUT_OR_PACKET_ERROR; // Error in cancel sequence
			}
		case ABORT1:
		case ABORT2: // Abort sequence
			return PACKET_RECEIVED; // Indicate abort
		default:
			return TIMEOUT_OR_PACKET_ERROR; // Invalid packet type
	}

	*data = c; // Store packet type

	// Receive sequence number and its complement
	for (i = 1; i < (packet_size + PACKET_OVERHEAD); i++)
	{
		if (Receive_Byte(data + i, timeout) != 0)
		{
			return TIMEOUT_OR_PACKET_ERROR; // Timeout or error
		}
	}
	
	// Calculate the length of the received packet
	if (data[PACKET_SEQNO_INDEX] != ((data[PACKET_SEQNO_COMP_INDEX] ^ 0xFF) & 0xFF))
	{
		return TIMEOUT_OR_PACKET_ERROR; // Sequence number mismatch
	}
	
	computed_crc = Cal_CRC16(&data[PACKET_HEADER], (uint32_t)packet_size);

	// Check CRC-16
	if (computed_crc != (uint16_t)((data[packet_size + 3] << 8) | data[packet_size + 4]))
	{
		return TIMEOUT_OR_PACKET_ERROR; // CRC mismatch
	}

	*length = packet_size; // Set the length of the received packet
	return PACKET_RECEIVED; // Packet received successfully
}

/**
 * @brief 接收 Ymodem 数据
 * 
 * @param buff 接收数据的缓冲区
 * @return int32_t 接收的数据大小, 0 : 成功, -1 : 错误, -2 : Flash 写入错误, -3 : 用户中止
 */
int32_t Ymodem_Receive(uint8_t *buff)
{
	uint8_t packet_data[PACKET_1K_SIZE + PACKET_OVERHEAD];
	uint8_t file_size[FILE_SIZE_LENGTH];
	uint8_t *file_ptr = NULL;
	uint8_t *buffer_ptr = NULL;

	int32_t packet_length = 0;
	int32_t session_done = 0;
	int32_t file_done = 0;
	int32_t packets_received = 0;
	int32_t errors;
	int32_t session_begin = 0;
	uint32_t size = 0;
	uint32_t flash_destination = APPLICATION_ADDRESS; // Example flash destination address
	uint32_t ram_source; // Example RAM source address

	for (session_done = 0, errors = 0, session_begin = 0; ;)
	{
		for (packets_received = 0, file_done = 0, buffer_ptr = buff; ;)
		{
			switch (Receive_Packet(packet_data, &packet_length, NAK_TIMEOUT))
			{
				case PACKET_RECEIVED: // Packet received successfully
					errors = 0;
					switch (packet_length)
					{
						case -1: // Abort by user
							Send_Byte(ACK);
							return 0;
						case 0: // End of Transmission
							Send_Byte(ACK);
							file_done = 1; // Indicate file done
							break;
						default: // Normal packet
							if ((packet_data[PACKET_SEQNO_INDEX] & 0xff) != (packets_received & 0xff))
							{
								Send_Byte(NAK);
							}
							else
							{
								if (packets_received == 0)
								{
									// Filename packet
									if (packet_data[PACKET_HEADER] != 0)
									{
										uint16_t i;
										// Filename packet has vaild data
										for (i = 0, file_ptr = packet_data + PACKET_HEADER; (*file_ptr != 0) && (i < FILE_NAME_LENGTH);)
										{
											file_name[i++] = *file_ptr++;
										}
										file_name[i++] = '\0'; // Null-terminate the filename
										for (i = 0, file_ptr++; (*file_ptr != ' ') && (i < FILE_SIZE_LENGTH - 1);)
										{
											file_size[i++] = *file_ptr++;
										}
										file_size[i++] = '\0'; // Null-terminate the file size
										Str2Int(file_size, &size); // Convert file size from string to integer

										if (size > (USER_FLASH_SIZE +1))
										{
											// End sessio
											Send_Byte(CAN); 
											Send_Byte(CAN);
											return -1; // Indicate error due to file size
										}

										// erase user application area
										FLASH_If_Erase(APPLICATION_ADDRESS);
										Send_Byte(ACK); // Acknowledge filename packet
										Send_Byte(CRC16); // Send CRC-16 request
									}
									// Filename packet is empty, end session
									else
									{
										Send_Byte(ACK); // Send cancel
										file_done = 1;
										session_done = 1; // End session
										break;
									}
								}
								// Data packet
								else
								{
									memcpy(buffer_ptr, packet_data + PACKET_HEADER, (uint32_t)packet_length);
									ram_source = (uint32_t)buff;

									// Write received data in Flash
									if (FLASH_If_Write(&flash_destination, (uint32_t *)ram_source, (uint16_t)packet_length / 4) == 0)
									{
										Send_Byte(ACK);
									}
									else // An error occurred during flash write
									{
										// End session
										Send_Byte(CAN); 
										Send_Byte(CAN);
										return -2;
									}
								}
								packets_received++;
								session_begin = 1;
							}
					}
					break;
				case ABORT_BY_USER:
					Send_Byte(CAN); // Abort
					Send_Byte(CAN);
					return -3; // Indicate abort
				default:
					if (session_begin > 0)
					{
						errors++;
					}
					if (errors > MAX_ERRORS)
					{
						Send_Byte(CAN); // Send cancel
						Send_Byte(CAN);
						return 0; // Indicate session end due to errors
					}
					Send_Byte(CRC16);
					break;
			}
			if (file_done != 0)
			{
				break;
			}
		}
		if (session_done != 0)
		{
			break; // Session done
		}
	}
	return (int32_t)size; // Return the size of the received data
}


//int32_t Ymodem_CheckResponse(uint8_t c)
//{
//	return 0;
//}

/**
 * @brief 准备起始帧
 * 
 * @param data 数据包缓冲区
 * @param f_name 文件名
 * @param length 文件大小
 */
static void Ymodem_PrepareIntialPacket(uint8_t *data, const uint8_t* f_name, uint16_t *length)
{
	uint16_t i, j;
	uint8_t file_ptr[10];

	// Prepare the initial packet with file name and size
	data[0] = SOH; 		// Start of 128 Byte data packet
	data[1] = 0x00; 	// Sequence number
	data[2] = 0xFF; 	// Complement of sequence number

	/******************** File name start ********************/
	// Check if filename packet has vaild data
	for(i = 0; (f_name[i] != '\0') && (i < FILE_NAME_LENGTH); i++)
	{
		data[PACKET_HEADER + i] = f_name[i];
	}
	data[PACKET_HEADER + i] = 0x00;	// Null-terminate the filename
	/******************** File name end ********************/


	/******************** File size start ********************/
	Int2Str(file_ptr, *length); // Convert file size to string
	for (j = 0, i = i + PACKET_HEADER + 1; file_ptr[j] != '\0' ;)
	{
		data[i++] = file_ptr[j++];
	}
	/******************** File size end ********************/

	// Fill the rest of the packet with zeros
	for (j = i; j < PACKET_SIZE + PACKET_HEADER; j++)
	{
		data[j] = 0x00;
	}
}

/**
 * @brief 准备数据帧
 * 
 * @param source_buff 源数据缓冲区
 * @param data 数据包缓冲区
 * @param pk_no 数据包序列号
 * @param size_blk 数据块大小
 */
static void Ymodem_PreparePacket(uint8_t *source_buff, uint8_t *data,
						  uint8_t pk_no, uint16_t size_blk)
{
	uint16_t i;
	uint16_t size;
	uint16_t packet_size;
	uint8_t * file_ptr = NULL;

	// Determine packet size
	packet_size = size_blk >= PACKET_1K_SIZE ? PACKET_1K_SIZE : PACKET_SIZE;

	size = size_blk < packet_size ? size_blk : packet_size;
	if (packet_size == PACKET_1K_SIZE)
	{
		data[0] = STX; // Start of 1K Byte data packet
	}
	else
	{
		data[0] = SOH; // Start of 128 Byte data packet
	}

	data[1] = pk_no; // Packet sequence number
	data[2] = ~pk_no; // Complement of sequence number
	file_ptr = source_buff;

	for (i = PACKET_HEADER; i < (size + PACKET_HEADER); i++)
	{
		data[i] = *file_ptr++; // Copy data from source buffer to packet
	}
	if (size <= packet_size)
	{
		for (i = size + PACKET_HEADER; i < (packet_size + PACKET_HEADER); i++)
		{
			// data[i] = 0x1a; // Fill remaining space with 0x1A (substitute character)
			data[i] = 0x00; // Fill remaining space with 0x00
		}
	}
}

/**
 * @brief 发送数据包
 * 
 * @param data 数据包缓冲区
 * @param length 数据包长度
 */
static void Ymodem_SendPacket(uint8_t *data, uint16_t length)
{
	for (uint16_t i = 0; i < length; i++)
	{
		Send_Byte(data[i]); // Send each byte of the packet
	}
}

/**
 * @brief 通过 Ymodem 协议发送一个文件
 * 
 * @param buff 数据缓冲区
 * @param f_name 文件名
 * @param f_size 文件大小
 * @return uint8_t 错误码, 0 : 成功, 非0 : 错误
 */
uint8_t Ymodem_Transmit(uint8_t *buff, const uint8_t* f_name, uint16_t f_size)
{
	uint8_t packet_data[PACKET_1K_SIZE + PACKET_OVERHEAD];
	uint8_t send_file_name[FILE_NAME_LENGTH];
	uint8_t *buff_ptr = NULL;
	uint8_t temp_checksum = 0;
	uint16_t temp_CRC = 0;
	uint8_t blk_number = 0;
	uint8_t receivedC[2];
	uint8_t CRC16_F = 1;		// CRC-16 标志位, 1-CRC16校验, 0-校验和
	uint8_t i;
	uint8_t errors = 0;
	uint32_t ack_received = 0;	// ACK 接收标志
	uint16_t size = 0;
	uint16_t pkt_size;			// 数据包大小, 128 Byte 或 1K Byte

	// Initialize file name buffer
	for (i = 0; i < (FILE_NAME_LENGTH - 1); i++)
	{
		send_file_name[i] = f_name[i];
	}


	/******************** Send first packet [START] *********************/
	// Prepare the initial packet with file name and size
	Ymodem_PrepareIntialPacket(packet_data, send_file_name, &f_size);

	while ((0 == ack_received) && (errors < 0x0A))
	{
		/* Send first packet */
		Ymodem_SendPacket(packet_data, PACKET_SIZE + PACKET_HEADER);

		if (CRC16_F)
		{
			temp_CRC = Cal_CRC16(&packet_data[PACKET_HEADER], PACKET_SIZE);
			Send_Byte((temp_CRC >> 8) & 0xFF); // Send high byte of CRC
			Send_Byte(temp_CRC & 0xFF); // Send low byte of CRC
		}
		else
		{
			temp_checksum = CalChecksum(&packet_data[PACKET_HEADER], PACKET_SIZE);
			Send_Byte(temp_checksum); // Send checksum
		}

		// Wait for ACK or NAK response
		if (Receive_Byte(receivedC, NAK_TIMEOUT) == 0)
		{
			if (receivedC[0] == ACK)
			{
				ack_received = 1; // Acknowledge received
			}
			else
			{
				errors++;
			}
		}
	}
	
	if (errors >= 0x0A)
	{
		return errors;
	}
	/******************** Send first packet [END] *********************/


	/******************** Send data packet [START] *********************/
	buff_ptr = buff;
	size = f_size;
	blk_number = 1;

	// 1024 bytes package is used to send the packets
	while (size > 0)
	{
		// Prepare next packet
		Ymodem_PreparePacket(buff_ptr, packet_data, blk_number, size);
		ack_received = 0;
		receivedC[0] = 0;
		errors = 0;

		while (0 == ack_received && (errors < 0x0A))
		{
			/* Send next packet */
			if (size >= PACKET_1K_SIZE)
			{
				pkt_size = PACKET_1K_SIZE; // Use 1K packet size
			}
			else
			{
				pkt_size = PACKET_SIZE; // Use 128 Byte packet size
			}
			Ymodem_SendPacket(packet_data, pkt_size + PACKET_HEADER);

			// Send CRC or checksum
			if (CRC16_F)
			{
				temp_CRC = Cal_CRC16(&packet_data[PACKET_HEADER], pkt_size);
				Send_Byte((temp_CRC >> 8) & 0xFF); // Send high byte of CRC
				Send_Byte(temp_CRC & 0xFF); // Send low byte of CRC
			}
			else
			{
				temp_checksum = CalChecksum(&packet_data[PACKET_HEADER], pkt_size);
				Send_Byte(temp_checksum); // Send checksum
			}

			// Wait for ACK or NAK response
			if (Receive_Byte(receivedC, NAK_TIMEOUT) == 0)
			{
				if (receivedC[0] == ACK)
				{
					ack_received = 1; // Acknowledge received
					if (size > pkt_size)
					{
						buff_ptr += pkt_size; // Move buffer pointer forward
						size -= pkt_size; // Decrease remaining size
						if (blk_number == (USER_FLASH_SIZE / 1024))
						{
							return 0xFF; // error
						}
						else
						{
							blk_number++;
						}
					}
					else
					{
						buff_ptr += pkt_size; 
						size = 0;
					}
				}
			}
			else
			{
				errors++; // Increment error count if no response
			}
		}
		
		if (errors >= 0x0A)
		{
			return errors; // Return error if too many attempts
		}
	}
	/******************** Send data packet [END] *********************/

	ack_received = 0;
	receivedC[0] = 0;
	receivedC[1] = 0;
	errors = 0;

	while (0 == ack_received && (errors < 0x0A))
	{
		Send_Byte(EOT); // Send End of Transmission

		receivedC[0] = (uint8_t)USART_ReceiveData(USART1); // Wait for ACK or NAK
		if (receivedC[0] == ACK)
		{
			ack_received = 1;
		}
		else
		{
			errors++; // Increment error count if no ACK
		}
		USART_ClearFlag(USART1, USART_FLAG_ORE); // Clear Overrun Error flag
	}
	if (errors >= 0x0A)
	{
		return errors; // Return error if too many attempts
	}
	

	// Last packet preparation
	ack_received = 0;
	receivedC[0] = 0;
	receivedC[1] = 0;
	errors = 0;
	packet_data[0] = SOH; // Start of 128 Byte data packet
	packet_data[1] = 0;
	packet_data[2] = 0xFF;
	for (i = PACKET_HEADER; i < PACKET_SIZE + PACKET_HEADER; i++)
	{
		packet_data[i] = 0x00; // Fill the rest of the packet with zeros
	}

	while (0 == ack_received && (errors < 0x0A))
	{
		/* Send packet */
		Ymodem_SendPacket(packet_data, PACKET_SIZE + PACKET_HEADER);

		temp_CRC = Cal_CRC16(&packet_data[PACKET_HEADER], PACKET_SIZE);
		Send_Byte((temp_CRC >> 8) & 0xFF); // Send high byte of CRC
		Send_Byte(temp_CRC & 0xFF); // Send low byte of CRC

		if (Receive_Byte(receivedC, NAK_TIMEOUT) == 0)
		{
			if (receivedC[0] == ACK)
			{
				ack_received = 1; // Acknowledge received
			}
		}
		else
		{
			errors++; // Increment error count if no response
		}
	}

	if (errors >= 0x0A)
	{
		return errors; // Return error if too many attempts
	}

	receivedC[0] = 0;
	while ((0 == ack_received) && (errors < 0x0A))
	{
		Send_Byte(EOT); // Send End of Transmission
		if ((Receive_Byte(receivedC, NAK_TIMEOUT) == 0) && (receivedC[0] == ACK))
		{
			ack_received = 1; // Acknowledge received
		}
		else
		{
			errors++; // Increment error count if no ACK
		}
		USART_ClearFlag(USART1, USART_FLAG_ORE); // Clear Overrun Error flag
	}
	if (errors >= 0x0A)
	{
		return errors; // Return error if too many attempts
	}
	
	return 0; // File transmission successful
}
